import os
import cv2
import argparse
from PIL import Image
import torch
from transformers import AutoProcessor, AutoModelForCausalLM, CLIPProcessor, CLIPModel

class VideoFrameAnalyzer:
    def __init__(self):
        print("正在加载图像分析模型...")
        # 加载CLIP模型用于图像特征提取和关键词生成
        self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
        
        # 加载图像描述生成模型（可选用更轻量级的模型）
        self.image_processor = AutoProcessor.from_pretrained("microsoft/git-base")
        self.image_model = AutoModelForCausalLM.from_pretrained("microsoft/git-base")
        
        # 预定义的关键词类别，用于匹配和筛选
        self.keyword_categories = [
            # 场景细分
            "室内家居", "办公室", "教室", "商场", "酒店", "医院", "餐厅", "公园", 
            "街道", "高速公路", "田野", "草原", "花园", "海滩", "沙漠", "森林",
            "山顶", "山谷", "湖泊", "河流", "瀑布", "海洋", "城市天际线", "乡村",
            
            # 物体细分
            "桌子", "椅子", "沙发", "床", "柜子", "电视", "电脑", "手机",
            "汽车", "自行车", "摩托车", "船", "飞机", "火车", "巴士", "卡车",
            "书籍", "花卉", "树木", "水果", "蔬菜", "肉类", "面包", "蛋糕",
        
        ]
        print("模型加载完成！")

    def extract_first_frame(self, video_path):
        """提取视频的第一帧"""
        video = cv2.VideoCapture(video_path)
        success, frame = video.read()
        video.release()
        
        if not success:
            raise Exception(f"无法从 {video_path} 读取视频帧")
        
        # 将OpenCV的BGR格式转换为RGB
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pil_image = Image.fromarray(frame_rgb)
        return pil_image

    def analyze_image(self, image):
        """分析图像并生成关键词"""
        # 使用CLIP模型提取图像特征
        inputs = self.clip_processor(
            text=self.keyword_categories,
            images=image,
            return_tensors="pt",
            padding=True
        )
        
        outputs = self.clip_model(**inputs)
        logits_per_image = outputs.logits_per_image
        probs = logits_per_image.softmax(dim=1)
        
        # 获取概率最高的8个关键词
        top_probs, top_indices = torch.topk(probs[0], k=8)
        top_keywords = [self.keyword_categories[idx] for idx in top_indices]
        
        # 使用图像描述模型生成更详细的描述（可选）
        pixel_values = self.image_processor(images=image, return_tensors="pt").pixel_values
        generated_ids = self.image_model.generate(
            pixel_values=pixel_values,
            max_length=50
        )
        description = self.image_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
        
        return top_keywords, description

    def rename_file(self, file_path, keywords):
        """将关键词添加到文件名中"""
        dir_path, filename = os.path.split(file_path)
        name, ext = os.path.splitext(filename)
        
        # 将关键词用下划线连接
        keywords_str = "_".join(keywords)
        new_name = f"{name}_{keywords_str}{ext}"
        new_path = os.path.join(dir_path, new_name)
        
        # 重命名文件
        os.rename(file_path, new_path)
        return new_path

    def process_video(self, video_path):
        """处理单个视频"""
        try:
            print(f"正在处理: {video_path}")
            # 提取第一帧
            first_frame = self.extract_first_frame(video_path)
            
            # 分析图像
            keywords, description = self.analyze_image(first_frame)
            print(f"图像描述: {description}")
            print(f"关键词: {', '.join(keywords)}")
            
            # 重命名文件
            new_path = self.rename_file(video_path, keywords)
            print(f"文件已重命名为: {new_path}")
            
            return True, new_path
        except Exception as e:
            print(f"处理 {video_path} 时出错: {str(e)}")
            return False, None

    def process_directory(self, directory_path, extensions=['.mp4', '.avi', '.mov', '.mkv']):
        """处理目录中的所有视频文件"""
        success_count = 0
        fail_count = 0
        
        for root, _, files in os.walk(directory_path):
            for file in files:
                if any(file.lower().endswith(ext) for ext in extensions):
                    file_path = os.path.join(root, file)
                    success, _ = self.process_video(file_path)
                    
                    if success:
                        success_count += 1
                    else:
                        fail_count += 1
        
        print(f"\n处理完成! 成功: {success_count}, 失败: {fail_count}")

def main():
    parser = argparse.ArgumentParser(description='视频首帧分析与重命名工具')
    parser.add_argument('path', help='视频文件路径或目录路径')
    parser.add_argument('--recursive', '-r', action='store_true', help='是否递归处理子目录')
    
    args = parser.parse_args()
    
    analyzer = VideoFrameAnalyzer()
    
    if os.path.isfile(args.path):
        analyzer.process_video(args.path)
    elif os.path.isdir(args.path):
        if args.recursive:
            analyzer.process_directory(args.path)
        else:
            # 只处理当前目录下的视频文件
            extensions = ['.mp4', '.avi', '.mov', '.mkv']
            for file in os.listdir(args.path):
                file_path = os.path.join(args.path, file)
                if os.path.isfile(file_path) and any(file.lower().endswith(ext) for ext in extensions):
                    analyzer.process_video(file_path)
    else:
        print(f"错误: 路径 '{args.path}' 不存在")

if __name__ == "__main__":
    main()
